use "../Common/CG.sml";

(* A structure which contains the relevant identifiers and functions of the Naive algorithm implementation, 
   adapted from the DICE 2010 paper. The algorithm pseudo-code can be found in "naive_algorithm.docx" *)
structure NaiveAnalysis =
struct


structure analysisKind = 
	struct 
		val nonlinear_Dvalue = D.Two;
		val name = "POLYNOMIAL";
	end; 
	
(* a global identifier which is updated (at the beginning of the algorithm execution) 
   to hold the number of variables in the analysed program. *)
val numVars = ref 0;
(* an iterator for debug print-out which is increased on each analysis interation over a loop body. *)
val itCount = ref 0;
(* debug mode -  when true - each algorithm stage will be printed, along with its output (and sometimes input). *)
val debug = CG.debug;	

(* prints s to the stdOut if debug is true. *)
fun debugPrint(s) = if (!debug) then print(s) else ();

(* returns flow i-(1)->i if i is not in source context and i <> x *)
fun computeFlowFromItoI(i : int, x : int, srcContext : CG.Context) =
	if (i <> x andalso (not (CG.inContext(i,srcContext)))) then [(i,i,D.One)] else [];

(* computes context graph for commands of the form "X := 0" *)
fun computeAsgnZeroGraph(x : int, preContexts : CG.Context list) : CG.graph =
	let 
		fun computeDataFlows(srcContext : CG.Context) : (int * int * D.t) list =
			let
				fun computeMultipleFlows(0, _) = []
                  |	computeMultipleFlows(i : int, srcContext : CG.Context) =
						computeFlowFromItoI(i,x,srcContext) @ computeMultipleFlows(i-1, srcContext);
			in
				computeMultipleFlows((!numVars), srcContext)
			end;
			
		fun computeEdge(srcContext : CG.Context) = 
			let
			    (* create target context identical to source context, at least at first... *)
				val tgtContext = CG.copyContext(!numVars, srcContext);
				(* include x in target context (equivalent to assignment of X := 0) *)
				val () = CG.addVarToContext(x,tgtContext);
				val m = mat.M(computeDataFlows(srcContext));
			in
				CG.E ( srcContext, tgtContext, ref (sop_bar.dfr(m, mat.r(m))) )
			end;			
	in
		CG.newGraph( map computeEdge preContexts )
	end;	

(* computes context graph for commands of the form "X := Y" *)
fun computeAsgnVarGraph(x : int, y : int, preContexts : CG.Context list) : CG.graph =
	let 
		fun computeDataFlows(srcContext : CG.Context) : (int * int * D.t) list =
			let
				fun computeFlowFromYtoX(i : int, srcContext : CG.Context)  =
					if (i = y andalso (not (CG.inContext(y, srcContext)))) then 
						[(y,x,D.One)]
					else 
						[];

				
				fun computeMultipleFlows(0, _) = []
                  |	computeMultipleFlows(i : int, srcContext : CG.Context) =
						let
							val flowYtoX = computeFlowFromYtoX(i,srcContext);
							val flowItoI = computeFlowFromItoI(i,x,srcContext);
						in
							flowYtoX @ flowItoI @ computeMultipleFlows(i-1, srcContext)
						end;
			in
				computeMultipleFlows((!numVars), srcContext)
			end;
			
		fun computeEdge(srcContext : CG.Context) = 
			let
			    (* create target context identical to source context, at least at first... *)
				val tgtContext = CG.copyContext(!numVars, srcContext);
				val y_in_srcContext = CG.inContext(y, srcContext);
				(* add x to target context if y is in source context (equivalent to assignment of X := 0),
				    otherwise exclude x from target context (equivalent to assignment of X := Y, where Y <> 0) *)
				val () = CG.setVarInContext(x, tgtContext, y_in_srcContext); 
				val m = mat.M(computeDataFlows(srcContext));
			in
				CG.E ( srcContext, tgtContext, ref (sop_bar.dfr(m, [])) )
			end;
		val g = CG.newGraph(map computeEdge preContexts);
	in
		g
	end;	

(* computes context graph for commands of the form "X := Y * Z" *)
fun computeAsgnTimesGraph(x, y, z, preContexts) : CG.graph =
	let 
		fun computeDataFlows(srcContext : CG.Context) : (int * int * D.t) list =
			let
				fun computeSingleFlow(i : int, srcContext : CG.Context) : (int * int * D.t) list =
					let 
						(* add flow y-(2)->x (or z-(2)->x) if y and z are both not in the source context 
						   (equivalent to assignment of X := Y*Z where Y>0 and Z>0 ) *)
						fun computeFlowFromYandZtoX(i : int, srcContext : CG.Context) =
							if ( (i = y orelse i = z) andalso
							     (not (CG.inContext(y, srcContext) orelse CG.inContext(z, srcContext))) )
							then [(i,x,D.Two)]
							else [];
					in
						computeFlowFromYandZtoX(i,srcContext) @ computeFlowFromItoI(i,x,srcContext)
					end;
				
				fun computeMultipleFlows(0,_) = []
                  |	computeMultipleFlows(i : int,  srcContext : CG.Context) =
						computeSingleFlow(i,srcContext) @ computeMultipleFlows(i-1,srcContext);
			in
				computeMultipleFlows((!numVars),srcContext)
			end;
			
		fun computeEdge(srcContext : CG.Context) = 
			let
			    (* create target context identical to source context, at least at first... *)
				val tgtContext = CG.copyContext(!numVars, srcContext);
				val y_or_z_in_srcContext = (CG.inContext(y, srcContext)) orelse (CG.inContext(z, srcContext));
				(* add x to target context if  y or z are in the source context (equivalent to assignment of X := 0),
				    otherwise exclude x from target context (equivalent to assignment of X := Y * Z, where Y>0 and Z > 0) *)
				val () = CG.setVarInContext(x, tgtContext, y_or_z_in_srcContext); 
				val mat = mat.M(computeDataFlows(srcContext));
			in
				CG.E ( srcContext, tgtContext, ref (sop_bar.dfr(mat, [])) )
			end;
	in
		CG.newGraph( map computeEdge preContexts )
	end;	

(* computes context graph for commands of the form "X := Y + Z" *)	
fun computeAsgnPlusGraph(x, y, z, preContexts) : CG.graph =
	let 
		fun computeDataFlows(srcContext : CG.Context) : (int * int * D.t) list =
			let
				fun computeSingleFlow(i : int, srcContext : CG.Context) : (int * int * D.t) list =
					let 
						(* add flow y-(1')->x (or z-(1')->x) if y (or z) are not in the source context 
						  (equivalent to assignment of X := Y+Z where Y>0 or Z>0 ) *)
						fun computeFlowFromYandZtoX(i : int, srcContext : CG.Context) =
							let
								val iNotInSrcContext = not (CG.inContext(i,srcContext));
								val yInSrcContext = CG.inContext(y, srcContext);
								val zInSrcContext = CG.inContext(z, srcContext);
							in
								if (i = y orelse i = z) then
									if (iNotInSrcContext andalso (not (yInSrcContext andalso zInSrcContext))) then 
										(if ((not yInSrcContext) andalso (not zInSrcContext)) then
											[(i,x,D.One')]
										else 
											[(i,x,D.One)])
									else []
								else []
							end;
					in
						 computeFlowFromYandZtoX(i,srcContext) @ computeFlowFromItoI(i,x,srcContext)
					end;
				
				fun computeMultipleFlows(0,_) = []
                  |	computeMultipleFlows(i : int, srcContext : CG.Context) =
							computeSingleFlow(i,srcContext) @ computeMultipleFlows(i-1,srcContext);								  

			in
				computeMultipleFlows((!numVars), srcContext)
			end;
			
		fun computeEdge(srcContext : CG.Context) = 
			let
			    (* create target context identical to source context, at least at first... *)
				val tgtContext = CG.copyContext(!numVars, srcContext);
				val y_and_z_in_srcContext = (CG.inContext(y, srcContext)) andalso (CG.inContext(z, srcContext));
				(* add x to target context if  y and z are both in the source context (equivalent to assignment of X := 0),
				    otherwise exclude x from target context (equivalent to assignment of X := Y + Z, where Y>0 or Z > 0) *)
				val () = CG.setVarInContext(x, tgtContext, y_and_z_in_srcContext);  
				val mat = mat.M(computeDataFlows(srcContext));
			in
				CG.E ( srcContext, tgtContext, ref (sop_bar.dfr(mat, mat.r(mat))) )
			end;
	in
		CG.newGraph( map computeEdge preContexts )
	end;

(* updates DFRs on all edges of graph g with all possible data flows of the form (i->j,i'->j'). *)
fun computeDoubleFlows(g : CG.graph) =
	let
		fun updateDoubleFlowsOnEdge(CG.E (p, q, dfr as (ref (sop_bar.dfr(m,_)))) ) =
			let		
				fun includeFlow (p : CG.Context,q : CG.Context) ((i,j),(i',j')) = 
								(not (CG.inContext(i,p))) andalso (not (CG.inContext(i',p))) andalso
								(not (CG.inContext(j,q))) andalso (not (CG.inContext(j',q))) andalso
								(i <> i' orelse j <> j');

				val r' = List.filter (includeFlow(p,q)) (mat.r(m)); 				
			in
				dfr := sop_bar.dfr(m,r')		
			end;

	in
		List.app updateDoubleFlowsOnEdge (CG.edges(g))
	end;

(* computes the context graph which corresponds to the assignment of expression e to the variable x with given pre-contexts. *)
fun computeAssignmentGraph (x : int, e : Core.Expr, preContexts : CG.Context list) : CG.graph =
	let
		fun computeSingleFlows (x : int, e : Core.Expr, preContexts : CG.Context list) : CG.graph =
			case e of
					  Core.Zero =>  						computeAsgnZeroGraph(x, preContexts)
					| Core.Var y =>  						computeAsgnVarGraph(x, y, preContexts)
					| Core.times(Core.Var y, Core.Var z) => computeAsgnTimesGraph(x, y, z, preContexts)									  
					| Core.plus(Core.Var y , Core.Var z) => computeAsgnPlusGraph(x, y, z, preContexts);
		
		val g' = computeSingleFlows(x, e, preContexts);
		val () = computeDoubleFlows(g');
	in
	   g'
	end;          

(* a global ientifier which holds a list of all contexts - for ease of reference anywhere in the algorithm. *)
val allContexts = ref ([CG.makeEmptyContext(1)]);

(* the algorithm entry point - gets as input a command c and a list of pre-contexts and 
   returns the context-graph which corresponds to the command c. *)
fun naive(c : Core.Cmd, preContexts : CG.Context list) : CG.graph = 
	let	
		(* Applies the loop correction operator on the context-graph g, according to loop rule L2 from the DICE paper. *)
		fun doLoopCorrection(l : int, g : CG.graph) : CG.graph =
			let
				(* overload the * operator for DFRs for ease of expressiveness of the DFR composition (in this function only). *)
				fun ((d1 :sop_bar.DFR) * (d2 :sop_bar.DFR)) = sop_bar.prod(d1,d2);
				
				(* applies the loop correction operator on a DFR *)
				fun LC(l,sop_bar.dfr(m,r)) =
					let
						val m' = mat.LC (l, analysisKind.nonlinear_Dvalue) m;
					in
						sop_bar.dfr(m', sop_bar.filterR(r,m'))
					end;
				
				(* computes LC on all edges (p,q') s.t. (p,q), (q,q') and (q,q) exist on graph. 
				   The result of this function is *)
				fun computeSingleEdge( _, []) = []
				  | computeSingleEdge(e as (CG.E (p,q,d)), (CG.E(p',q',d'))::t) = 
						(if (CG.eqContexts(q,p')) then
							(case CG.getEdgeTag(g,CG.UE (q,q)) of
											SOME dfr => [(CG.E (p,q', ref ( !d * LC(l, !dfr) * !d' ) ))] 
														   
										  | NONE 	 => nil)
						else nil) @ computeSingleEdge(e,t);
							
				fun computeMultipleEdges([],_) = []
				  | computeMultipleEdges(e::t, fullEdgeList) = 
							computeSingleEdge(e,fullEdgeList) @ computeMultipleEdges(t,fullEdgeList);	
				
				val sortedEdges = CG.sortEdges(CG.edges(g));
				
				val () = CG.addEdges(g, computeMultipleEdges(sortedEdges, sortedEdges));
			in
				g
			end;
		
		(* computes a single iteration of a loop whose body is represented by the graph g (after some m >= 0 iterations). *)
		fun doSingleIteration (g : CG.graph ref) g'  =
			let
				val () = debugPrint("================================\nIteration #" ^ Int.toString(!itCount) ^ ":\ncurrent graph(g):\n");
				val () = debugPrint(CG.prtEdges(!g));
				val () = debugPrint("graph to compose(g'):\n");
				val () = debugPrint(CG.prtEdges(g'));
				val g'' = CG.compose(!g,g');
				val () = debugPrint("composed graph(g''):\n" ^ CG.prtEdges(g''));
				val () = (g := CG.join(!g, g''));
				val () = debugPrint("join result graph(g):\n");
				val () = debugPrint(CG.prtEdges(!g));
				val () = debugPrint("================================\n");
				val () = itCount := !itCount + 1;
				val _ = if (!debug) then TextIO.inputLine(TextIO.stdIn) else NONE;
			in
				!g
			end;

		(* computes the result graph of any m > 0 iterations, where the input graph is the graph representing the loop-body and 
           the possibility of 0 iterations (skip). This function corresponds to loop rule L1 from the DICE paper. *)		
		fun doMIterations(x : int, g : CG.graph ref) : CG.graph =
			let
				fun lfp f (g : CG.graph) = 
					let
						val fg = (f g);
					in
						if (CG.equal(fg,g)) then g else lfp f fg
				end;
				
				val g' = CG.clone(!g);
				val () = itCount := 0;
			in
				lfp (doSingleIteration g) g'
			end
		
		(* Computes the result graph of a loop whose body is represented by g and the iterating variable is x. *)
		fun doLoop(x : int, g : CG.graph ref) : CG.graph =
			let
				(* the context graph g is joined with a graph which corresponds to the Skip command with the same pre-contexts.	
				   This action corresponds to applying loop rule L0 from the DICE paper. *)
				val () = g := CG.join(!g, naive(Core.Skip, CG.preContexts(!g))) 
			in
				doLoopCorrection(x, doMIterations(x, g))
			end;
			
		fun recurse(c : Core.Cmd, preContexts : CG.Context list) =
			case c of Core.Seq(c1,c2)      => 
												let
													val g1 = naive(c1, preContexts);
													val g2 = naive(c2, CG.postContexts(g1));
												in
													CG.compose(g1, g2)
												end
					|  Core.Choice(c1,c2)  => CG.join(naive(c1, preContexts), naive(c2, preContexts))
					|  Core.Skip		   => recurse(Core.Asgn(1, Core.Var 1), preContexts)
					|  Core.Asgn(i, e)     => computeAssignmentGraph(i, e, case preContexts of 
																			  [] => !allContexts
																		    | _  => preContexts)
					|  Core.Loop(i, c)     => doLoop(i, ref (naive(c, preContexts)));
					
		val () = debugPrint("starting to analyse command:\n");
		val () = debugPrint(Core.prt(c));
		val () = debugPrint("Pre-contexts: " ^ CG.prtContexts(preContexts) ^ "\n");
		val g = recurse(c, preContexts);
		val () = debugPrint("Finished analysing command:\n");
		val () = debugPrint(Core.prt(c));
		val () = debugPrint("graph:\n");
		val () = debugPrint(CG.prtEdges(g));
		val () = debugPrint("============================================\n");
	in
		g
	end;

(* the function to use for invoking the analysis algorithm on a command c. *)
fun analyse(c : Core.Cmd) : CG.graph = 
	let
		(* initialize "global" identifiers. *)
		val () = (numVars := Core.highest_var(c));
		val () = allContexts := CG.allContexts(!numVars);
	in
		naive(c,[])
	end;
	
end; (* structure NaiveAnalysis *)



